package com.zl.learning.arithmetic.rbtree;

import lombok.Data;
import lombok.Getter;

public class RBTree {

    private static final boolean RED = false;
    private static final boolean BLACK = true;

    @Getter
    private Node root;
    private int size;


    public void preOrderTraverse() {
        preOrderTraverse(root);
    }

    private void preOrderTraverse(Node p) {
        if (p == null) return;
        System.out.print(p.data + " ");
        preOrderTraverse(p.left);
        preOrderTraverse(p.right);
    }

    public void inOrderTraverse() {
        inOrderTraverse(root);
    }

    private void inOrderTraverse(Node p) {
        if (p == null) return;
        inOrderTraverse(p.left);
        System.out.print(p.data + " ");
        inOrderTraverse(p.right);
    }

    public void postOrderTraverse() {
        postOrderTraverse(root);
    }

    private void postOrderTraverse(Node p) {
        if (p == null) return;
        postOrderTraverse(p.left);
        postOrderTraverse(p.right);
        System.out.print(p.data + " ");
    }

    public void add(Integer data) {
        Node t = root;
        if (t == null) {
            root = new Node(data, null);
            size = 1;
            return;
        }
        Node parent;
        boolean cmp;
        do {
            parent = t;
            cmp = t.data > data;
            if (cmp) {
                t = t.left;
            } else {
                t = t.right;
            }
        } while (t != null);
        Node e = new Node(data, parent);
        if (cmp) {
            parent.left = e;
        } else {
            parent.right = e;
        }
        fixAfterInsertion(e);
        size++;
    }

    public void delete(Integer data) {
        Node p = getNode(data);
        delete(p);
    }

    private void delete(Node p) {
        size--;
        //如果p即有左孩子也有右孩子
        if (p.left != null && p.right != null) {
            Node s = successor(p);
            p.data = s.data;
            p = s;
        }
        Node replacement = (p.left != null ? p.left : p.right);
        if (replacement != null) {
            replacement.parent = p.parent;
            if (p.parent == null) {
                root = replacement;
            } else if (p == p.parent.left) {
                p.parent.left = replacement;
            } else {
                p.parent.right = replacement;
            }
            p.left = p.right = p.parent = null;
            if (p.color == BLACK) {
                fixAfterDeletion(replacement);
            }
        } else if (p.parent == null) {
            root = null;
        } else {
            if (p.color == BLACK) {
                fixAfterDeletion(p);
            }
            if (p == p.parent.left) {
                p.parent.left = null;
            } else if (p == p.parent.right) {
                p.parent.right = null;
            }
            p.parent = null;
        }
    }

    private void fixAfterDeletion(Node x) {
        while (x != root && colorOf(x) == BLACK) {
            if (x == leftOf(parentOf(x))) {
                Node sib = rightOf(parentOf(x));
                if (colorOf(sib) == RED) {
                    setColor(sib, BLACK);
                    setColor(parentOf(x), RED);
                    rotateLeft(parentOf(x));
                    sib = rightOf(parentOf(x));
                }
                if (colorOf(leftOf(sib)) == BLACK && colorOf(rightOf(sib)) == BLACK) {
                    setColor(sib, RED);
                    x = parentOf(x);
                } else {
                    if (colorOf(rightOf(sib)) == BLACK) {
                        setColor(leftOf(sib), BLACK);
                        setColor(sib, RED);
                        rotateRight(sib);
                        sib = rightOf(parentOf(x));
                    }
                    setColor(sib, colorOf(parentOf(x)));
                    setColor(parentOf(x), BLACK);
                    setColor(rightOf(sib), BLACK);
                    rotateLeft(parentOf(x));
                    x = root;
                }
            } else {
                Node sib = leftOf(parentOf(x));
                if (colorOf(sib) == RED) {
                    setColor(sib, BLACK);
                    setColor(parentOf(x), RED);
                    rotateRight(parentOf(x));
                    sib = leftOf(parentOf(x));
                }
                if (colorOf(rightOf(sib)) == BLACK && colorOf(leftOf(sib)) == BLACK) {
                    setColor(leftOf(sib), RED);
                    x = parentOf(x);
                } else {
                    if (colorOf(leftOf(sib)) == BLACK) {
                        setColor(rightOf(sib), BLACK);
                        setColor(sib, RED);
                        rotateLeft(sib);
                        sib = leftOf(parentOf(x));
                    }
                    setColor(sib, colorOf(parentOf(x)));
                    setColor(parentOf(x), BLACK);
                    setColor(leftOf(sib), BLACK);
                    rotateRight(parentOf(x));
                    x = root;
                }
            }
        }
        setColor(x, BLACK);
    }

    private Node successor(Node t) {
        if (t == null) {
            return null;
        } else if (t.right != null) {
            Node p = t.right;
            while (p.left != null) {
                p = p.left;
            }
            return p;
        } else {
            Node p = t.parent;
            Node ch = t;
            while (p != null && ch == p.right) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

    private Node getNode(Integer data) {
        Node p = root;
        while (p != null) {
            int cmp = data - p.data;
            if (cmp < 0) {
                p = p.left;
            } else if (cmp > 0) {
                p = p.right;
            } else {
                return p;
            }
        }
        return null;
    }

    private Node parentOf(Node p) {
        return (p == null ? null : p.parent);
    }

    private Node leftOf(Node p) {
        return p == null ? null : p.left;
    }

    private Node rightOf(Node p) {
        return p == null ? null : p.left;
    }

    private boolean colorOf(Node p) {
        return p == null ? BLACK : p.color;
    }

    private void setColor(Node p, boolean c) {
        if (p != null) p.color = c;
    }

    private void fixAfterInsertion(Node x) {
        x.color = RED;
        while (x != null && x != root && colorOf(parentOf(x)) == RED) {//X不是null并且x不是根节点并且x的父节点是红色
            if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {//如果父节点是祖父节点的左孩子
                Node y = rightOf(parentOf(parentOf(x)));//叔叔节点是右孩子
                if (colorOf(y) == RED) {//如果叔叔节点是红色
                    setColor(parentOf(x), BLACK); //把父节点和叔叔节点变黑、把祖父变红
                    setColor(y, BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    x = parentOf(parentOf(x));
                } else {//如果叔叔是黑色
                    if (x == rightOf(parentOf(x))) { //x是父亲的右节点、以父亲为当前节点，左旋
                        x = parentOf(x);
                        rotateLeft(x);
                    }
                    setColor(parentOf(x), BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    rotateRight(parentOf(parentOf(x)));
                }
            } else {//父节点是祖父的右孩子
                Node y = leftOf(parentOf(parentOf(x)));
                if (colorOf(y) == RED) {//叔叔是红色 把父亲和叔叔变黑，把祖父变红 把祖父变成当前节点
                    setColor(parentOf(x), BLACK);
                    setColor(y, BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    x = parentOf(parentOf(x));
                } else {//叔叔是黑色
                    if (x == leftOf(parentOf(x))) { //X 是父亲的左孩子 把父亲变为当前节点 右旋
                        x = parentOf(x);
                        rotateRight(x);
                    }
                    setColor(parentOf(x), BLACK); //
                    setColor(parentOf(parentOf(x)), RED);
                    rotateLeft(parentOf(parentOf(x)));
                }
            }
        }
        root.color = BLACK;
    }

    /**
     * 左旋 把p变成 p 的右子树的左子树
     *
     * @param p
     */
    private void rotateLeft(Node p) {
        if (p != null) {
            Node r = p.right;
            //把p右子树的左子树变成p的右子树
            p.right = r.left;
            if (r.left != null) r.left.parent = p;
            r.parent = p.parent;
            if (p.parent == null) root = r;
            else if (p.parent.left == p) p.parent.left = r;
            else p.parent.right = r;
            r.left = p;
            p.parent = r;
        }
    }

    /**
     * 右旋
     *
     * @param p
     */
    private void rotateRight(Node p) {
        if (p != null) {
            Node l = p.left;
            p.left = l.right;
            if (l.right != null) l.right.parent = p;
            l.parent = p.parent;
            if (p.right == null) root = l;
            else if (p.parent.right == p) p.parent.right = l;
            else p.parent.left = l;
            l.right = p;
            p.parent = l;
        }
    }

    @Override
    public String toString() {
        return "rbtree";
    }

    @Data
    static class Node {
        private transient Node parent;
        private Node right;
        private Node left;
        private int data;
        private boolean color;

        public Node() {
        }

        public Node(int data, Node parent) {
            this.data = data;
            this.parent = parent;
        }

        @Override
        public String toString() {
            return this.data + "--" + (this.color ? "黑色" : "红色");
        }

    }
}
